from utils.utils import *
from torch import nn, Tensor
from torch.nn import functional as F

class StyleLoss(nn.Module):
    """图片风格损失函数

    通过特征映射的 gram 矩阵进行计算，该损失函数是为了保证在生成的图像中保风格图片的样式。
    ..ref: https://github.com/eriklindernoren/Fast-Neural-Style-Transfer/blob/master/train.py
    """

    def __init__(self, weight):
        super().__init__()
        self.weight = weight

    def forward(self, batch_y: Tensor, batch_gram_y_hat: Tensor) -> Tensor:
        style_loss = 0.0
        for y, gram_y_hat in zip(batch_y, batch_gram_y_hat):
            gram_y = gram_matrix(y)
            style_loss += F.mse_loss(gram_y, gram_y_hat[:batch_y[0].shape[0], :, :])
        return self.weight * style_loss
